package hn.cch.spring_boot_security_rest.wrapper;

import org.springframework.util.StreamUtils;

import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.*;
import java.util.Collections;
import java.util.Enumeration;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class ApiRequestWrapper extends HttpServletRequestWrapper {

    private Map<String, String> headerMap = new ConcurrentHashMap<>();


    public void addHeader(String name, String header) {
        this.headerMap.put(name, header);
    }

    @Override
    public String getHeader(String name) {
        String header = super.getHeader(name);
        return header != null ? header : this.headerMap.get(name);
    }

    @Override
    public Enumeration<String> getHeaderNames() {
        List<String> names = Collections.list(super.getHeaderNames());
        names.addAll(this.headerMap.keySet());
        return Collections.enumeration(names);
    }

    @Override
    public Enumeration<String> getHeaders(String name) {
        List<String> headerList = Collections.list(super.getHeaders(name));
        if (this.headerMap.containsKey(name)) {
            headerList.add(this.headerMap.get(name));
        }
        return Collections.enumeration(headerList);
    }

    @Override
    public String getQueryString() {
        return super.getQueryString();
    }


    private byte[] body;

    private ByteArrayInputStream bais;

    private ServletInputStream sis;

    private BufferedReader br;


    public ApiRequestWrapper(HttpServletRequest request) throws IOException {
        super(request);
        byte[] body = StreamUtils.copyToByteArray(request.getInputStream());
        bais = new ByteArrayInputStream(body);
        sis = new ApiServletInputStream(bais);
        br = new BufferedReader(new InputStreamReader(sis));
    }

    @Override
    public ServletInputStream getInputStream() throws IOException {
        return sis;
    }

    @Override
    public BufferedReader getReader() throws IOException {
        return br;
    }


    public byte[] getRequestBody() {
        return body;
    }
}
